{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8b13b5b",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from pydrake.geometry.optimization import (\n",
    "    GraphOfConvexSetsOptions,\n",
    "    HPolyhedron,\n",
    "    Point,\n",
    ")\n",
    "from pydrake.planning import GcsTrajectoryOptimization\n",
    "from scipy.spatial import ConvexHull"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2da43708",
   "metadata": {},
   "source": [
    "# Environment setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8582fff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sr2 = np.sqrt(2)\n",
    "\n",
    "# obstacles\n",
    "obstacles = [\n",
    "    np.array(\n",
    "        [\n",
    "            [3.4, 2.6],\n",
    "            [3.4, 4.6],\n",
    "            [2.4, 4.6],\n",
    "            [2.4, 2.6],\n",
    "            [1.4, 2.2],\n",
    "            [3.8, 0.2],\n",
    "            [4.8, 1.2],\n",
    "        ]\n",
    "    ),\n",
    "    np.array([[1.4, 2.8], [2.2, 2.8], [2.2, 4.6], [1.4, 4.6]]),\n",
    "    np.array([[1.0, 2.6], [1.0, 5.0], [0.4, 5.0], [0.4, 2.6]]),\n",
    "    np.array([[1.0, 2.4], [1.0, 0.0], [0.4, 0.0], [0.4, 2.4]]),\n",
    "    np.array([[3.8, 3.0], [3.8, 5.0], [4.4, 5.0], [4.4, 3.0]]),\n",
    "    np.array([[3.8, 2.8], [3.8, 2.6], [5.0, 2.6], [5.0, 2.8]]),\n",
    "]\n",
    "\n",
    "# vertices of the safe regions\n",
    "vertices = [\n",
    "    np.array([[0.4, 0.0], [0.4, 5.0], [0.0, 5.0], [0.0, 0.0]]),\n",
    "    np.array([[0.4, 2.4], [1.0, 2.4], [1.0, 2.6], [0.4, 2.6]]),\n",
    "    np.array([[1.4, 2.2], [1.4, 4.6], [1.0, 4.6], [1.0, 2.2]]),\n",
    "    np.array([[1.4, 2.2], [2.4, 2.6], [2.4, 2.8], [1.4, 2.8]]),\n",
    "    np.array([[2.2, 2.8], [2.4, 2.8], [2.4, 4.6], [2.2, 4.6]]),\n",
    "    np.array([[1.4, 2.2], [1.0, 2.2], [1.0, 0.0], [3.8, 0.0], [3.8, 0.2]]),\n",
    "    np.array([[3.8, 4.6], [3.8, 5.0], [1.0, 5.0], [1.0, 4.6]]),\n",
    "    np.array([[5.0, 0.0], [5.0, 1.2], [4.8, 1.2], [3.8, 0.2], [3.8, 0.0]]),\n",
    "    np.array([[3.4, 2.6], [4.8, 1.2], [5.0, 1.2], [5.0, 2.6]]),\n",
    "    np.array([[3.4, 2.6], [3.8, 2.6], [3.8, 4.6], [3.4, 4.6]]),\n",
    "    np.array([[3.8, 2.8], [4.4, 2.8], [4.4, 3.0], [3.8, 3.0]]),\n",
    "    np.array([[5.0, 2.8], [5.0, 5.0], [4.4, 5.0], [4.4, 2.8]]),\n",
    "]\n",
    "\n",
    "x_min = np.min(np.vstack(vertices), axis=0)\n",
    "x_max = np.max(np.vstack(vertices), axis=0)\n",
    "\n",
    "x_start = np.array([0.2, 0.2])\n",
    "x_goal = np.array([4.8, 4.8])\n",
    "\n",
    "\n",
    "def make_hpolytope(V):\n",
    "    ch = ConvexHull(V)\n",
    "    return HPolyhedron(ch.equations[:, :-1], -ch.equations[:, -1])\n",
    "\n",
    "\n",
    "regions = [make_hpolytope(V) for V in vertices]\n",
    "\n",
    "\n",
    "def environment_setup():\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.axis(\"square\")\n",
    "\n",
    "    plt.xlim([x_min[0], x_max[0]])\n",
    "    plt.ylim([x_min[1], x_max[1]])\n",
    "\n",
    "    tick_gap = 0.2\n",
    "\n",
    "    def n_ticks(x_min, x_max):\n",
    "        return round((x_max - x_min) / tick_gap) + 1\n",
    "\n",
    "    x_ticks = np.linspace(x_min[0], x_max[0], n_ticks(x_min[0], x_max[0]))\n",
    "    y_ticks = np.linspace(x_min[1], x_max[1], n_ticks(x_min[1], x_max[1]))\n",
    "    plt.xticks(x_ticks)\n",
    "    plt.yticks(y_ticks)\n",
    "\n",
    "    label_gap = 0.5\n",
    "\n",
    "    def keep_label(t):\n",
    "        return np.isclose(t % label_gap, 0) or np.isclose(\n",
    "            t % label_gap, label_gap\n",
    "        )\n",
    "\n",
    "    x_labels = [int(t) if keep_label(t) else \"\" for t in x_ticks]\n",
    "    y_labels = [int(t) if keep_label(t) else \"\" for t in y_ticks]\n",
    "    plt.gca().set_xticklabels(x_labels)\n",
    "    plt.gca().set_yticklabels(y_labels)\n",
    "\n",
    "    plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c9ec31",
   "metadata": {},
   "outputs": [],
   "source": [
    "environment_setup()\n",
    "\n",
    "for O in obstacles:\n",
    "    plt.fill(*O.T, fc=\"lightcoral\", ec=\"k\", zorder=4)\n",
    "\n",
    "plt.plot(*x_start, \"kx\")\n",
    "plt.plot(*x_goal, \"kx\")\n",
    "\n",
    "plt.text(0.2, 0.35, \"$q_0$\", ha=\"center\", va=\"bottom\")\n",
    "plt.text(4.8, 4.65, \"$q_T$\", ha=\"center\", va=\"top\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87afe7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "environment_setup()\n",
    "\n",
    "for V in vertices:\n",
    "    plt.fill(*V.T, fc=\"lightcyan\", ec=\"k\", zorder=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0520b584",
   "metadata": {},
   "source": [
    "# Minimum-distance problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d1c11d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_trajectory(traj):\n",
    "    plt.figure(figsize=(10, 10))\n",
    "\n",
    "    for O in obstacles:\n",
    "        plt.fill(*O.T, fc=\"lightcoral\", ec=\"k\", zorder=4)\n",
    "\n",
    "    plt.plot(*traj.value(traj.start_time()), \"kx\")\n",
    "    plt.plot(*traj.value(traj.end_time()), \"kx\")\n",
    "    times = np.linspace(traj.start_time(), traj.end_time(), 1000)\n",
    "    waypoints = traj.vector_values(times)\n",
    "    plt.plot(*waypoints, \"b\", zorder=5)\n",
    "\n",
    "    plt.axis(\"square\")\n",
    "    plt.xlim([x_min[0], x_max[0]])\n",
    "    plt.ylim([x_min[1], x_max[1]])\n",
    "    plt.xticks(range(6))\n",
    "    plt.yticks(range(6))\n",
    "    plt.grid(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05f59a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "trajopt = GcsTrajectoryOptimization(2)\n",
    "gcs_regions = trajopt.AddRegions(regions, order=1)\n",
    "source = trajopt.AddRegions([Point(x_start)], order=0)\n",
    "target = trajopt.AddRegions([Point(x_goal)], order=0)\n",
    "trajopt.AddEdges(source, gcs_regions)\n",
    "trajopt.AddEdges(gcs_regions, target)\n",
    "trajopt.AddPathLengthCost()\n",
    "options = GraphOfConvexSetsOptions()\n",
    "options.convex_relaxation = True\n",
    "[traj, result] = trajopt.SolvePath(source, target, options)\n",
    "plot_trajectory(traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2c4a76",
   "metadata": {},
   "source": [
    "# Minimum-time problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e727dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "qdot_min = -1\n",
    "qdot_max = 1\n",
    "\n",
    "\n",
    "def solve_min_time(order, continuity_order):\n",
    "    trajopt = GcsTrajectoryOptimization(2)\n",
    "    gcs_regions = trajopt.AddRegions(regions, order=order)\n",
    "    source = trajopt.AddRegions([Point(x_start)], order=0)\n",
    "    target = trajopt.AddRegions([Point(x_goal)], order=0)\n",
    "    trajopt.AddEdges(source, gcs_regions)\n",
    "    trajopt.AddEdges(gcs_regions, target)\n",
    "    trajopt.AddTimeCost()\n",
    "    trajopt.AddVelocityBounds([qdot_min] * 2, [qdot_max] * 2)\n",
    "    for o in range(1, continuity_order + 1):\n",
    "        print(f\"adding C{o} constraints\")\n",
    "        trajopt.AddPathContinuityConstraints(o)\n",
    "    options = GraphOfConvexSetsOptions()\n",
    "    options.convex_relaxation = True\n",
    "    [traj, result] = trajopt.SolvePath(source, target, options)\n",
    "\n",
    "    return traj, result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3b8f7a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_velocity(traj):\n",
    "    vel = traj.MakeDerivative()\n",
    "\n",
    "    plt.figure(figsize=(10, 8))\n",
    "\n",
    "    for i in range(vel.get_number_of_segments()):\n",
    "        v = vel.segment(i)\n",
    "        times = np.linspace(v.start_time(), v.end_time(), 500)\n",
    "        values = v.vector_values(times)\n",
    "        plt.plot(times, values[0], color=\"tab:blue\")\n",
    "        plt.plot(times, values[1], color=\"tab:orange\")\n",
    "\n",
    "    plt.xlim([traj.start_time(), traj.end_time()])\n",
    "    plt.xticks(np.arange(int(np.ceil(traj.end_time() / 2))) * 2)\n",
    "    plt.yticks(np.linspace(qdot_min, qdot_max, 5))\n",
    "    plt.xlabel(\"Time $t$\")\n",
    "    plt.ylabel(\"Velocity $\\dot{q}$\")\n",
    "    plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba49e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "order = 1\n",
    "continuity = 0\n",
    "traj, results = solve_min_time(order, continuity)\n",
    "\n",
    "plot_trajectory(traj)\n",
    "\n",
    "plot_velocity(traj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b578c00f",
   "metadata": {},
   "outputs": [],
   "source": [
    "order = 6\n",
    "continuity = 1\n",
    "# regularizer = [1e-1, 1e-1]\n",
    "traj, results = solve_min_time(order, continuity)\n",
    "\n",
    "plot_trajectory(traj)\n",
    "\n",
    "plot_velocity(traj)\n",
    "\n",
    "path = GcsTrajectoryOptimization.NormalizeSegmentTimes(traj)\n",
    "plot_velocity(path)\n",
    "plt.ylabel(\"Path Velocity $\\dot{r}$\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
